# coding = utf-8

'''
对Lits数据集的读取和分析
'''

import nibabel as nib
import os
import SimpleITK as sitk
import numpy as np

def read_data():

  size_100w = []
  size_10w = []
  size_1w = []
  size_1k = []
  size_small = []
  size_zero = []

  lits_root_path = "F:\Dataset\Liver\LITS"
  for i in range(131):
      mask_file = os.path.join(lits_root_path, "segmentation-{}.nii".format(i))
      image_file = os.path.join(lits_root_path, "volume-{}.nii".format(i))

      image = sitk.GetArrayFromImage(sitk.ReadImage(image_file))
      mask = sitk.GetArrayFromImage(sitk.ReadImage(mask_file))

      temp = np.zeros(mask.shape)
      temp[mask == 2] = 1

      sum_size = int(temp.sum())

      if sum_size >= 1000000:
          size_100w.append(i)
      elif sum_size >= 100000:
          size_10w.append(i)
      elif sum_size >= 10000:
          size_1w.append(i)
      elif sum_size >= 1000:
          size_1k.append(i)
      elif sum_size >= 1:
          size_small.append(i)
      else:
          size_zero.append(i)

  print(len(size_100w), size_100w)
  print(len(size_10w), size_10w)
  print(len(size_1w), size_1w)
  print(len(size_1k), size_1k)
  print(len(size_small), size_small)
  print(len(size_zero), size_zero)





if __name__ == '__main__':
    read_data()